from cryptography.fernet import Fernet
from .cipher import encrypt, decrypt
from .taichu_protect import *
import builtins
import os
import sys
import io
from loguru import logger as logging
# from env import MODEL_PATH

logging.info(f'MODEL_PATH: {os.getenv("MODEL_PATH")}')

def need_hook(file_path: str) -> bool:
    if file_path.startswith(os.getenv("MODEL_PATH")) and not (file_path.endswith('.py') or file_path.endswith('tokenizer.model') or file_path.endswith('.pth') or file_path.endswith('.pt') \
                                                              or file_path.endswith('qformer_tokenizer/tokenizer_config.json') or file_path.endswith('qformer_tokenizer/tokenizer.json')):
        logging.info(f'need_hook file: {file_path}')
        sys.stdout.flush()
        return True
    return False


# 备份原始的 open 函数
_org_open = builtins.open


class EncryptedFile:
    def __init__(self, file, mode, *args, **kwargs):
        # 创建新的文件对象来避免递归的调用问题
        self.file = io.BytesIO()
        if 'r' in mode:
            cur_mode = mode
            if 'b' not in mode:
                cur_mode = f'{mode}b' 
            with _org_open(file, cur_mode) as f:
                if file.endswith('.json') or file.endswith('.txt'):
                    data = f.read()
                    decrypted = decrypt(data, cur_mode)
                    self.file.write(decrypted)
                else:
                    while True:
                        data = f.read(block_size)
                        if len(data) < block_size:
                            self.file.write(data)
                            break
                        decrypted = aes_de(data)
                        self.file.write(decrypted)
        self.file.seek(0)
        self.mode = mode
        self.filename = file
        self.writed = False

    def __getattr__(self, name):
        try:
            if name == 'mode':
                return self.mode
            return getattr(self.file, name)
        except Exception as e:
            logging.info(f'getattr error: {e}')            

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def read(self, size=-1):
        # 在读取文件时进行解密操作
        data = self.file.read(size)
        if 'b' not in self.mode:
            data = data.decode('utf-8')
        return data

    def readinto(self, b):
        # 在读取文件时进行解密操作
        size = len(b)
        data = self.file.read(size)
        b[:size] = data[:size]
        return len(data)

    def readline(self, size=-1):
        # 在读取文件时进行解密操作
        data = self.file.readline(size)
        if 'b' not in self.mode:
            data = data.decode('utf-8')
        return data

    def readlines(self, hint=-1):
        # 在读取文件的行时进行解密操作
        lines = self.file.readlines(hint)
        if 'b' not in self.mode:
            lines = [line.decode('utf-8') for line in lines]
        return lines

    def write(self, data):
        # 在写入文件时进行加密操作
        if not isinstance(data, bytes):
            data = data.encode('utf-8')
        # logging.info(f'write data: {data}')
        return self.file.write(data)

    def writelines(self, lines):
        # 在写入文件时进行加密操作
        for line in lines:
            self.write(line)

    def flush(self):
        return self.file.flush()

    def seek(self, offset, whence=os.SEEK_SET):
        return self.file.seek(offset, whence)

    def tell(self):
        return self.file.tell()

    def peek(self, size):
        cur = self.tell()
        data = self.file.read(size)
        self.file.seek(cur)
        return data

    def truncate(self, size=None):
        return self.file.truncate(size)

    def close(self):
        self._write()
        return self.file.close()

    def _write(self):
        if 'w' in self.mode and not self.writed:
            with _org_open(self.filename, 'wb') as f:
                raw_data = self.file.getvalue()
                file_size = len(raw_data)
                # logging.info(f'raw_data: {raw_data}')
                if self.filename.endswith('.json'):
                    encrypted = encrypt(raw_data, self.mode)
                    f.write(encrypted)
                else:
                    offset = 0
                    while True:
                        end_pos = offset + block_size
                        if end_pos >= file_size:
                            end_pos = file_size
                            data = raw_data[offset:end_pos]
                            f.write(data)
                            break
                        data = raw_data[offset:end_pos]
                        encrypted = aes_en(data)
                        f.write(encrypted)
                        offset = end_pos
            self.writed = True

    def __del__(self):
        # 关闭文件对象
        self._write()
        self.file.close()

def encrypted_open(filename, mode="r", *args, **kwargs):
    # 通过原始的 open 函数创建文件对象
    filename = str(filename)
    if need_hook(filename):
        # logging.info(f'need_hook file: {filename}')
        # 返回加密后的文件对象
        return EncryptedFile(filename, mode, *args, **kwargs)
    else:
        return _org_open(filename, mode, *args, **kwargs)

# 重写 open 函数，用 encrypted_open 替换原始的 open
def install():
    builtins.open = encrypted_open
    pass

def uninstall():
    builtins.open = _org_open


if __name__ == "__main__":
    install()
    #filename = 'test.json'
    filename = './config.json'
    #uninstall()
    #with open(filename, 'w') as f:
        #line1 = "hello world, man, you are so great\n".encode('utf-8')
        #line2 = "so beautiful".encode('utf-8')
    #    line1 = "hello world, man, you are so great\n"
    #    line2 = "so beautiful"
    #    lines = [line1, line2]
        # f.write(line1)
        # f.write(line2)
    #    f.writelines(lines)

    with open(filename, 'r') as f:
        content = f.read()
        print(content)
